import cv2
import rosbag
from sensor_msgs.msg import Image
from cv_bridge import CvBridge
import numpy as np


def extract_images_from_bag(bag_file, image_topic):
    images = []
    with rosbag.Bag(bag_file, "r") as bag:
        for topic, msg, t in bag.read_messages(topics=[image_topic]):
            bridge = CvBridge()
            cv_image = bridge.imgmsg_to_cv2(msg, desired_encoding="bgr8")
            images.append(cv_image)
    return images


def save_video(images, output_file, fps):
    height, width, _ = images[0].shape
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    video_writer = cv2.VideoWriter(output_file, fourcc, fps, (width, height))
    for image in images:
        video_writer.write(image)
    video_writer.release()


def main():
    bag_file = "/home/hw/yellow_2024-08-04-10-17-08.bag"
    image_topic = "/perception/zed2i/rgb_leftView"  # 修改为你的图像话题名称
    output_file = "only_yellow.mp4"
    fps = 30  # 视频帧率

    images = extract_images_from_bag(bag_file, image_topic)
    save_video(images, output_file, fps)


if __name__ == "__main__":
    main()
